import pickle
import numpy as np
import gradio as gr
from PIL import Image  # 导入 PIL 库中的 Image 模块

# 加载保存的KNN模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

# 定义预测函数
def predict_digit(image):
    # 调整大小为8x8
    image = Image.fromarray(image).convert('L').resize((8, 8))
    image = np.array(image) / 32.0  # 缩放到0-32之间
    # 将图像展平为一维数组并进行预测
    prediction = knn_model.predict(image.flatten().reshape(1, -1))
    return int(prediction[0])

# 创建Gradio界面
iface = gr.Interface(
    fn=predict_digit,  # 设置要调用的预测函数
    inputs=gr.Sketchpad(),  # 设置输入类型为手写板
    outputs=gr.Label(num_top_classes=1),  # 设置输出类型为标签，显示预测结果
    title="手写数字识别",  # 设置界面的标题
    description="在画板上绘制一个数字，模型将预测该数字。"  # 设置界面的描述
)

# 启动Gradio界面
iface.launch(share=True)